package cn.yanceysong.kafka;

import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.redisson.Redisson;
import org.redisson.api.RBucket;
import org.redisson.api.RedissonClient;
import org.redisson.client.codec.StringCodec;
import org.redisson.config.Config;

import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.TimeUnit;

public class RedisSink implements SinkFunction<String>, AutoCloseable {
    public static RedissonClient redissonClient;
    private String redisHost;
    private Integer port;

    public RedisSink(String redisHost, Integer port) {
        this.redisHost = redisHost;
        this.port = port;
    }

    @Override
    public void invoke(String value, Context context) throws Exception {
        if (RedisSink.redissonClient == null) {
            Config config = new Config();
            config.useSingleServer()
                    .setAddress("redis://" + redisHost + ":" + port)
                    .setPassword("123456")
                    .setDatabase(0);
            config.setCodec(new StringCodec());
            config.setNettyThreads(0);
            RedisSink.redissonClient = Redisson.create(config);
        }
        RBucket<String> bucket = RedisSink.redissonClient.getBucket("kafka:" + getRandomString(5));
        bucket.set(value, 10, TimeUnit.SECONDS);
    }

    @Override
    public void close() throws Exception {
        if (RedisSink.redissonClient != null) {
            RedisSink.redissonClient.shutdown();
        }
    }

    public static String getRandomString(int length) {
        String base = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < length; i++) {
            int number = ThreadLocalRandom.current().nextInt(base.length());
            sb.append(base.charAt(number));
        }
        return sb.toString();
    }
}
